import pickle
import numpy as np
import gradio as gr
from PIL import Image
import cv2

# 加载预训练的最优KNN模型
def load_best_knn_model(model_path='best_knn_model.pkl'):
    try:
        with open(model_path, 'rb') as f:
            model = pickle.load(f)
        print(f"成功加载模型：{model_path}")
        return model
    except FileNotFoundError:
        raise Exception(f"错误：未找到模型文件{model_path}，请先运行optimal_knn.py生成模型")

best_knn = load_best_knn_model()

# 优化的图像预处理与预测函数
def predict_digit(image):
    # 处理空输入
    if image is None:
        return None  # 空输入时返回None而非0
    
    # 从Gradio Sketchpad提取图像数据
    if isinstance(image, dict):
        image_array = image.get('image', None)
        if image_array is None:
            return None
        # 转换为PIL图像格式
        image = Image.fromarray(np.array(image_array, dtype=np.uint8))
    
    # 确保输入是PIL图像
    if not isinstance(image, Image.Image):
        return None
    
    try:
        # 1. 转换为灰度图
        gray_image = image.convert('L')
        
        # 2. 检查是否全为白色（空绘制）
        if np.array(gray_image).mean() > 250:  # 几乎全白
            return None
        
        # 3. 调整尺寸为8×8
        resized_image = gray_image.resize((8, 8), Image.Resampling.LANCZOS)
        
        # 4. 颜色反转（黑底白字 -> 白底黑字）
        inverted_image = Image.eval(resized_image, lambda x: 255 - x)
        
        # 5. 像素值归一化到0-16（匹配训练数据范围）
        np_image = np.array(inverted_image, dtype=np.float32)
        np_image = (np_image / 255) * 16
        
        # 6. 展平为64维特征向量
        flattened_image = np_image.flatten().reshape(1, -1)
        
        # 调试信息：打印预处理后的像素值（可选）
        # print("预处理后的数据:", flattened_image)
        
        # 7. 模型预测
        prediction = best_knn.predict(flattened_image)
        return int(prediction[0])
    
    except Exception as e:
        print(f"预处理错误：{str(e)}")
        return None  # 错误时返回None

# 创建Gradio交互界面
with gr.Blocks(title="KNN手写数字识别") as demo:
    gr.Markdown("# KNN手写数字识别")
    gr.Markdown("### 请在下方手写板绘制0-9的数字，点击提交查看预测结果")
    gr.Markdown("#### 提示：尽量写大且清晰，占据手写板中心区域")
    
    with gr.Row():
        # 手写板组件
        sketchpad = gr.Sketchpad(
            height=200,
            width=200,
            label="手写板"
        )
        # 预测结果显示
        result = gr.Number(
            label="预测结果",
            precision=0,
            scale=2,
            interactive=False
        )
    
    with gr.Row():
        submit_btn = gr.Button("提交预测", variant="primary")
        clear_btn = gr.Button("清除手写板")
    
    # 绑定按钮事件
    submit_btn.click(
        fn=predict_digit,
        inputs=sketchpad,
        outputs=result
    )
    clear_btn.click(
        fn=lambda: (None, None),  # 清除时结果也设为None
        inputs=[],
        outputs=[sketchpad, result]
    )

# 启动应用（修改端口号）
if __name__ == "__main__":
    print("启动Gradio应用...")
    print("访问链接: http://127.0.0.1:7862")  # 端口改为7862
    demo.launch(
        server_name="127.0.0.1",
        server_port=7862,  # 这里修改端口号
        share=False
    )